require "test/unit"

module Test
  module Unit
    class TestCase
      class << self
        alias_method :method_added_without_attributes, :method_added
        def method_added(name)
          method_added_without_attributes(name)
          if defined?(@current_attributes)
            attributes = {}
            kept_attributes = {}
            @current_attributes.each do |attribute_name, attribute|
              attributes[attribute_name] = attribute[:value]
              kept_attributes[attribute_name] = attribute if attribute[:keep]
            end
            set_attributes(name, attributes)
            @current_attributes = kept_attributes
          end
        end

        def attribute(name, value, options={}, *tests)
          unless options.is_a?(Hash)
            tests << options
            options = {}
          end
          @current_attributes ||= {}
          if tests.empty?
            @current_attributes[name] = options.merge(:value => value)
          else
            tests.each do |test|
              set_attribute(test, {name => value})
            end
          end
        end

        def bug(value, *tests)
          attribute(:bug, value, *tests)
        end

        def set_attributes(test_name, attributes)
          return if attributes.empty?
          test_name = normalize_test_name(test_name)
          @attributes ||= {}
          @attributes[test_name] ||= {}
          @attributes[test_name] = @attributes[test_name].merge(attributes)
        end

        def attributes(test_name)
          test_name = normalize_test_name(test_name)
          @attributes ||= {}
          @attributes[test_name]
        end

        private
        def normalize_test_name(test_name)
          "test_#{test_name.to_s.sub(/^test_/, '')}"
        end
      end

      alias_method :run_without_attributes, :run
      def run(result, &block)
        run_without_attributes(TestResultAttributesSupport.new(result, self),
                               &block)
      end

      def attributes
        self.class.attributes(@method_name) || {}
      end
    end

    class TestResultAttributesSupport
      def initialize(result, test)
        @result = result
        @test = test
      end

      def add_failure(failure)
        failure.attributes = @test.attributes
        method_missing(:add_failure, failure)
      end

      def add_error(error)
        error.attributes = @test.attributes
        method_missing(:add_error, error)
      end

      def method_missing(name, *args, &block)
        @result.send(name, *args, &block)
      end
    end

    module AttributesFormatter
      private
      def format_attributes
        return '' if attributes.empty?
        attributes.collect do |key, value|
          "  #{key}: #{value}"
        end.join("\n") + "\n"
      end
    end

    class Failure
      include AttributesFormatter

      attr_accessor :attributes

      alias_method :long_display_without_attributes, :long_display
      def long_display
        test_name_re = Regexp.escape(@test_name)
        long_display_without_attributes.sub(/(^#{test_name_re}.*\n)/,
                                            "\\1#{format_attributes}")
      end
    end

    class Error
      include AttributesFormatter

      attr_accessor :attributes

      alias_method :long_display_without_attributes, :long_display
      def long_display
	test_name_re = Regexp.escape(@test_name)
        long_display_without_attributes.sub(/(^#{test_name_re}:\n)/,
                                            "\\1#{format_attributes}")
      end
    end
  end
end
